package com.haidechizi.leetcode;

import java.util.HashMap;
import java.util.Map;
import java.util.TreeMap;

public class LFUCache {

    private Map<Integer, Integer> map;

    private TreeMap<Integer, Integer> treeMap;

    private int capacity;

    public LFUCache(int capacity) {
        this.capacity = capacity;
        map = new HashMap<>(capacity);
        treeMap = new TreeMap<>((k1,k2)-> {
            return treeMap.get(k1) - treeMap.get(k2);
        });

    }

    public int get(int key) {
        return map.getOrDefault(key, -1);
    }

    public void put(int key, int value) {
        if (map.containsKey(key)) {
            map.put(key, value);
            Integer orDefault = treeMap.get(key);
            orDefault++;
            treeMap.remove(key);
            treeMap.put(key, orDefault);
        } else {
            if (map.size() >= capacity) {
                int firstKey = treeMap.firstKey();
                map.remove(firstKey);
                treeMap.remove(firstKey);
            }
            map.put(key, value);
            treeMap.put(key, 1);
        }
    }
}
